// #pragma once
#include <cstdio>
#include <iostream>
#include "myType.cuh"
#include "cuda_runtime_api.h"
#include <sys/time.h>
#include <typeinfo>
#include <algorithm>
#include <chrono>

#define GPU_CHECK(call)                               \
do                                                    \
{                                                     \
    const cudaError_t error_code = call;              \
    if (error_code != cudaSuccess)                    \
    {                                                 \
        printf("CUDA Error:\n");                      \
        printf("    File:       %s\n", __FILE__);     \
        printf("    Line:       %d\n", __LINE__);     \
        printf("    Error code: %d\n", error_code);   \
        printf("    Error text: %s\n",                \
            cudaGetErrorString(error_code));          \
        exit(1);                                      \
    }                                                 \
} while (0)

inline __device__ void gpu_swap(iAS* a, iAS* b) {
    iAS temp = *b;
    *b = *a;
    *a = temp;
}

void cpu_swap(iAS** a, iAS** b) {
    iAS* temp = *a;
    *a = *b;
    *b = temp;
}

// 大于比较符号
inline __device__  bool compare(const iAS& a, const iAS& b) {
    if(a.score != b.score) return a.score > b.score;
    return a.idx < b.idx;
}

bool cpu_compare(const iAS& a, const iAS& b) {
    if(a.score != b.score) return a.score > b.score;
    return a.idx < b.idx;
}

__device__ void max_heapify(iAS *arr, const int left, const int right) {
    //建立父节点指标和子节点指标
    register int dad = left;
    register int son = dad * 2 + 1;
    while (son <= right) { //若子节点指标在范围内才做比较
        if (son + 1 <= right && compare(arr[son], arr[son + 1])) //先比较两个子节点大小，选择最小的
            son++;
        if (compare(arr[son], arr[dad])) //如果父节点小于子节点代表调整完毕，直接跳出函数
            return;
        else { //否则交换父子内容再继续子节点和孙节点比较
            // printf("%d,%d",dad,son);
            gpu_swap(&arr[dad], &arr[son]);
            dad = son;
            son = dad * 2 + 1;
        }
    }
}

// 构建大顶堆
__global__ void gpu_topk(const iAS *data, iAS *d_result, const int length, const bool need_sort) {
    register auto tid = blockIdx.x * blockDim.x + threadIdx.x;
    
    if(threadIdx.x >= TOPK || tid >= length) {
        return;
    }

    __shared__ iAS temp_result[TOPK];
    
    // 初始化前TOPK个
    temp_result[threadIdx.x].score = data[tid].score; 
    temp_result[threadIdx.x].idx = data[tid].idx;
    __syncthreads();

    // 0号线程完成后续工作
    if(threadIdx.x > 0) {
        return;
    }
    register int temp_result_length = (TOPK <= length - tid ? TOPK : length - tid);

    // 建堆
#pragma unroll
    for (register int j = temp_result_length / 2 - 1; j >= 0; j--) {
        max_heapify(temp_result, j, TOPK - 1);
    }
    
    // 插入TOPK后面的所有数字
#pragma unroll
    for (register int i = temp_result_length; i < blockDim.x && i + tid < length; i++) {
        // printf("%d\n",i + tid);
        if(compare(data[i + tid], temp_result[0])) {
            temp_result[0].score = data[i + tid].score;
            temp_result[0].idx = data[i + tid].idx;
            max_heapify(temp_result, 0, TOPK - 1);
        }
    }

    register int offset = blockIdx.x * TOPK;

    if(need_sort) {
#pragma unroll
        for(register auto i = temp_result_length - 1; i >= 0; i--) {
            // printf("%d %f\n", temp_result[0].idx, temp_result[0].score);
            
            d_result[i + offset].score = temp_result[0].score;
            d_result[i + offset].idx = temp_result[0].idx;
            gpu_swap(&temp_result[0], &temp_result[i]);
            max_heapify(temp_result, 0, i - 1);
        }
        return;
    }

    // 赋值
#pragma unroll
    for(register auto i = 0; i < temp_result_length; i++) {
        // printf("%d %f\n", temp_result[i].idx, temp_result[i].score);
        // printf("%d", i + blockIdx.x * TOPK);
        d_result[i + offset].score = temp_result[i].score;
        d_result[i + offset].idx = temp_result[i].idx;
        // printf("%d %f\n", d_result[i + blockIdx.x * TOPK].idx, d_result[i + blockIdx.x * TOPK].score);        
    }
}


// Initialize data on the host.
void initialize_data(iAS *dst, unsigned int nitems) {
    // Fixed seed for illustration
    srand(2047);

    // Fill dst with random values
    for (unsigned i = 0; i < nitems; i++) {
        dst[i].idx = i;
        dst[i].score = rand() % nitems;
    }
}

// Verify the results.
void print_results(iAS *results_d, int n) {
    iAS *results_h = new iAS[n];
    GPU_CHECK(cudaMemcpy(results_h, results_d, n * sizeof(iAS), cudaMemcpyDeviceToHost));
    std::cout << "Sort data : \n";
    for (int i = 0; i < n; ++i){
        std::cout << results_h[i].idx << " " << results_h[i].score << " " << "\n";
    }
    std::cout << std::endl;
    delete[] results_h;
}

void check_results(iAS *results_d, iAS *data_h, int n)
{
    iAS *results_h = new iAS[n];
    GPU_CHECK(cudaMemcpy(results_h, results_d, n*sizeof(iAS), cudaMemcpyDeviceToHost));

    for (int i = 0; i < n ; ++i) {
        if(results_h[i].idx != data_h[i].idx) {
            std::cout << "Invalid item[" << i << "]: " << results_h[i].idx << " greater than " << data_h[i].idx << std::endl;
            std::cout << "Invalid item[" << i << "]: " << results_h[i].score << " greater than " << data_h[i].score << std::endl;
            // exit(EXIT_FAILURE);
        }
    }
        // if (cpu_compare(results_h[i], results_h[i-1]))
        // {
        //     std::cout << "Invalid item[" << i-1 << "]: " << results_h[i-1].idx << " greater than " << results_h[i].idx << std::endl;
        //     std::cout << "Invalid item[" << i-1 << "]: " << results_h[i-1].score << " greater than " << results_h[i].score << std::endl;
        //     exit(EXIT_FAILURE);
        // }

    std::cout << "OK" << std::endl;
    delete[] results_h;
}

// Main entry point.
int main(int argc, char **argv) {
    // Find/set device and get device properties
    int device = 0;
    cudaDeviceProp deviceProp;
    GPU_CHECK(cudaGetDeviceProperties(&deviceProp, device));
    
    if (!(deviceProp.major > 3 ||
          (deviceProp.major == 3 && deviceProp.minor >= 5))) {
        printf("GPU %d - %s  does not support CUDA Dynamic Parallelism\n Exiting.",
            device, deviceProp.name);
        return 0;
    }
    

    int num_items = atoi(argv[1]); // 7853052
    bool verbose = false;


    // Create input data
    iAS *h_data = 0;
    iAS *d_data = 0;
    iAS *d_result = 0;
    iAS *d_mid_result = 0;

    // Allocate CPU memory and initialize data.
    std::cout << "Initializing data:" << std::endl;
    h_data = (iAS *)malloc(num_items * sizeof(iAS));
    initialize_data(h_data, num_items);
    if (verbose) {
       std::cout << "Raw  data : \n";
        for (int i = 0; i < num_items; i++)
            std::cout << h_data[i].idx << " " << h_data[i].score << "\n";
    }

    // Allocate GPU memory.
    GPU_CHECK(cudaMalloc((void **)&d_data, num_items * sizeof(iAS)));
    GPU_CHECK(cudaMalloc((void **)&d_result, num_items * sizeof(iAS)));
    GPU_CHECK(cudaMalloc((void **)&d_mid_result, num_items * sizeof(iAS)));
    
    // GPU_CHECK(cudaMemset(d_mid_result, -1, num_items * sizeof(iAS)));
    // GPU_CHECK(cudaMemset(d_result, -1, num_items * sizeof(iAS)));

    GPU_CHECK(cudaMemcpy(d_data, h_data, num_items * sizeof(iAS), cudaMemcpyHostToDevice));
    
    // Execute
    std::cout << "Running quicksort on " << num_items << " elements" << std::endl;
    

    std::chrono::high_resolution_clock::time_point t1 = std::chrono::high_resolution_clock::now();

    //grid要不大于10

    // 第一轮 grid = 7669
    int block = 512;
    int temp_grid = (num_items + block - 1) / block;
    int temp_length = num_items;
    gpu_topk<<<temp_grid, block>>>(d_data, d_mid_result, temp_length, false);
    cudaDeviceSynchronize();
    temp_length = (temp_grid - 1) * TOPK + (temp_length % block < TOPK ? temp_length % block : TOPK);
    // printf("a%d\n", temp_length);
    // print_results(d_mid_result, temp_length);


    while(temp_length > 1024) {
        temp_grid = (temp_length + block - 1) / block;
        cpu_swap(&d_mid_result, &d_result);
        // print_results(d_result, temp_length);
        gpu_topk<<<temp_grid, block>>>(d_result, d_mid_result, temp_length, false);
        cudaDeviceSynchronize();

        temp_length = (temp_grid - 1) * TOPK + (temp_length % block < TOPK ? temp_length % block : TOPK);
        // printf("b%d\n", temp_length);
        // print_results(d_mid_result, temp_length);
    }
    // printf("while finish\n");
   
    gpu_topk<<<1, temp_length>>>(d_mid_result, d_result, temp_length, true);
    cudaDeviceSynchronize();
    // print_results(d_result, TOPK);

    
    std::chrono::high_resolution_clock::time_point t2 = std::chrono::high_resolution_clock::now();


    std::cout<<std::chrono::duration_cast<std::chrono::milliseconds>(t2-t1).count()<<std::endl;

    // print result
    // print_results(d_result, TOPK);
    // check result
    // std::cout << "Validating results: ";
    std::sort(h_data, h_data + num_items);
    check_results(d_result, h_data, TOPK);
    free(h_data);
    GPU_CHECK(cudaFree(d_data));

    return 0;
}

/* 
编译：nvcc -o quicksort_cuda --gpu-architecture=sm_70 -rdc=truequicksort_cuda.cu
*/